#!/usr/bin/env bash

# set -x

# 灵感来自 https://github.com/worthable/openai-terminal-assistant

# Tip: 大部分大模型都会对一些符号增强识别

# 引号 "..." 单行包含目标字段
# 三引号 """...""" 多行包含目标字段
# 井号 “#” 和三个井号 “###” 与引号作用相当，常用于强调命令
# 分割线 --- / === 常用于输入命令或数据

# 拦截脚本内容执行
[[ "$0" == "$SHELL" ]] && echo "\e[31m错误:\e[0m 不允许以命令方式执行脚本内容" && return

########################################################################

# 逐字输出
print(){
  while IFS= read -r LINE;do
    local STR="$LINE"
    local LEN=${#STR}
    for ((i=0;i<LEN;i++));do
      echo -n "${STR:i:1}"
      sleep 0.001
    done
    echo
  done
}

########################################################################

# 帮助
chathelp(){
CHAT_HRLP_ALL=(
  '/help'   '显示此帮助'
  '/exit'   '退出会话'
  '/clear'  '清空对话'
  '/pass'   '空操作'
  '/files'  '上传文件'
  '/file'   '提交文件'
  '/exec'   '执行命令'
  '/eval'   '生成命令'
  '/set'    '设置功能'
  '/cls'    '清除屏幕'
)
CHAT_KEYS_ALL=(
  '^C'      '退出会话'
  '^L'      '清除屏幕'
  '^R'      '搜索记录'
)
cat <<EOF
$(printf "\e[35mOption:\e[0m")
$(printf "  \e[32m%-10s\e[0m%-4s\n" ${CHAT_HRLP_ALL[@]})

$(printf "\e[35mkeys:\e[0m")
$(printf "  \e[32m%-10s\e[0m%-4s\n" ${CHAT_KEYS_ALL[@]})

EOF
}

########################################################################

# 错误输出并退出
exmsg(){
  exec echo -e $2 "\e[31m错误:\e[0m $1" >&2
}

# 仅错误输出
emsg(){
  echo -e $2 "\e[31m错误:\e[0m $1" >&2
}

# 常时输出
msg(){
  echo -e $2 "\e[35m$arg0:\e[0m $1"
}

# 更具交互性的错误输出
erron(){
  if $ERRON;then
    printf "\e[1A\r\e[K"
    emsg "$1\e[?25l\e[8m" -n
    sleep 1
    printf "\r\e[?25h\e[28m\e[K"
  else
    emsg "$1" $2
  fi
}

# 调试输出
debug(){
  $DEBUG && echo -e $2 "\e[33mdebug:\e[0m $1" >&2
}

########################################################################

# 设置扩展
setload(){
  local alias value type name
  case "$#" in
    0)
      for ((i=0;i<${#SET[@]};i++));do
        alias=${SET_ALIAS[i]}
        value=${!SET[i]}
        printf "%-20s%-4s\n" $alias $value
      done
      ;;
    1)
      alias=$1
      for ((i=0;i<${#SET[@]};i++));do
        type=${SET_ALIAS[i]}
        value=${!SET[i]}
        if [[ "$alias" == "$type" ]];then
          echo "$value"
          return
        fi
      done
      erron "未知的选项 $alias"
      ;;
    2)
      alias=$1
      value=$2
      for ((i=0;i<${#SET[@]};i++));do
        type=${SET_ALIAS[i]}
        name=${SET[i]}
        if [[ "$alias" == "$type" ]];then
          if [[ "${value}" =~ ^(false|true)$ ]];then
            eval "$name=$value"
          else
            erron "仅允许使用布尔值"
          fi
          return
        fi
      done
      erron "未知的选项 $alias"
  esac
}

########################################################################

# 文件类型
filetype(){
  local typeof="$(file -b "${1:--}" 2>/dev/null)"
  [[ "$typeof" =~ ^(.+UTF-8.+|ASCII text|JSON text data)$ ]]
}

# 判断变量类型
typeval(){
  if [ -z "${!1}" ];then
    echo "zero"      # 空变量
  elif declare -p "$1" 2>/dev/null | grep -q 'declare -a';then
    echo "array"     # 数组
  elif [[ "${!1}" =~ ^-?[0-9]+(\.[0-9]+)?$ ]];then
    echo "number"    # 数值
  elif [[ "${!1}" =~ ^https?://[0-9a-zA-Z./?=\&%-_]+$ ]];then
    echo "url"       # 链接
  elif [[ "${!1}" =~ ^(false|true)$ ]];then
    echo "bollean"   # 布尔值
  elif [[ "${!1}" =~ ^.+$ ]];then
    echo "string"    # 字符串
  fi
}

########################################################################

# 等待阻塞任务
getwait(){
  # 使用 kill 命令判断进程存在，不存在即停止等待
  while $CHAT;do
    for i in '⣷' '⣯' '⣟' '⡿' '⢿' '⣻' '⣽' '⣾';do
      if kill -0 $PROCESS &>/dev/null;then
        echo -ne "\e[?25l\r\e[1;32m$i\e[0m 正在等待" >&2
        sleep 0.05
      else
        echo -ne "\e[?25h\r\e[K" >&2
        exit 0
      fi
    done
  done
}

########################################################################

# 判断响应错误
geterr(){
  if ! $(jq 'has("error")' <<< "$1");then
    ERROR_MSG=$(jq --exit-status -r '.error.message' <<< "$1" 2>/dev/null)
    API_STATUS=$?
    if ((API_STATUS==1)) && [[ "$ERROR_MSG" == "null" ]];then
      ERROR_MSG="该次请求未正确被接受"
    fi
    erron "$ERROR_MSG"
    return 1
  fi
}

########################################################################

# 管理文件
apifile(){
  if [ -z "$1" ];then
    emsg "未指定参数"
    return 1
  fi
  local fileapi="$OPENAI_BASE_URL/files"
  case "$1" in
    push)    # 提交文件
      if [ -z "$2" ];then
        erron "未指定参数"
        return 1
      fi
      if [ ! -f "$2" ];then
        erron "找不到文件 $2"
        return 1
      fi
      if [[ ! "$2" =~ ^(jsonl?|csv)$ ]];then
        erron "文件不符合要求"
        return 1
      fi
      local purpose="$3"
      if [[ "${purpose:=assistants}" =~ ^(assistants|fine-tune|file-extract|batch)$ ]];then
        CURL_FILE_GET="$(
          (
            curl -s -X POST "$fileapi" -H "$API_AUTH" \
              -F "purpose=$purpose" -F "file=@$2" 2>/dev/null ||
            printf '{"error": {"message": "文件提交时遇到错误"}}'
          ) &
          PROCESS=$!
          getwait
        )"
        geterr "$CURL_FILE_GET" &&
        msg "文件提交完成 $(jq -r '.id' <<< "$CURL_FILE_GET")"
      else
        erron "非法的用途 $purpose"
        return 1
      fi
      ;;
    list)    # 查看文件列表
      CURL_FILE_GET="$(
        (
          curl -s -X POST "$fileapi" -H "$API_AUTH" \
            -F "file=@$2" 2>/dev/null ||
          printf '{"error": {"message": "查找文件时遇到错误"}}'
        ) &
        PROCESS=$!
        getwait
      )"
      geterr "$CURL_FILE_GET"
      FILES_ID=($(jq -r '.data[].id' <<< "$CURL_FILE_GET"))
      FILES_NAME=($(jq -r '.data[].filename' <<< "$CURL_FILE_GET"))
      ;;
    rm)      # 删除文件
      FILES_LIST="${#FILES_ID[@]}"
      if ((FILES_LIST==0));then
        erron "无可删除的文件"
        return 1
      fi
      while true;do
        for ((i=0;i<FILES_LIST;i++));do
          printf "$i. %-30s%-2s\n" "${FILES_ID[i]}" "${FILES_NAME[i]}"
        done
        msg "选择要删除的文件: " -n
        read choose
        Choose="$(typeval choose)"
        max_choose="$[FILES_LIST-1]"
        if [[ "$Choose" == "number" ]];then
          if ((choose<0));then
            erron "选择不得为负"
          elif ((choose>max_choose));then
            erron "选择过大"
          else
            CURL_FILE_GET="$(
              (
                curl -s -X DELETE "$fileapi/${FILES_ID[choose]}" -H "$API_AUTH" 2>/dev/null ||
                printf '{"error": {"message": "文件删除时遇到错误"}}'
              ) &
              PROCESS=$!
              getwait
            )"
            geterr "$CURL_FILE_GET" && break
          fi
        elif [[ "$Choose" == "string" && "$choose" == "/exit" ]];then
          msg "取消选择"
          break          
        fi
      done
  esac
}

########################################################################

# 命令名称
arg0="$(basename ${BASH_SOURCE[0]} .sh)"

# 历史记录
HISTFILE=~/.bash_${arg0}_chat_history

# 解析参数
OPENAI_USER_PROMPT=""
OPENAI_SYSTEM_PROMPT=""
STDIN_INPUT=""

# 检查标准输入
if ! [ -t 0 ];then
  ! filetype - && exmsg "非法的文件类型"
  STDIN_INPUT="$(cat)"
  STDIN=true
else
  STDIN=false
fi

########################################################################

# 简单检查依赖命令
NOT_FOUND=()
for i in sed jq curl grep file;do
  command -v $i &>/dev/null || NOT_FOUND+=($i)
done
((${#NOT_FOUND[@]}==0)) || exmsg "缺失命令 ${NOT_FOUND[@]}"

########################################################################

# 解析参数
for arg in $@;do
  case "$1" in
    --config|-f)
      if [ -z $2 ];then
        exmsg "未指定文件"
      fi
      if [ -f $2 ];then
        if (
             ! grep -Ev '^\s*(export\s+)?[a-zA-Z_][a-zA-Z0-9_]*=.*(#.*)?$' "$2" |
             grep -qEv '^\s*$|^\s*#' &&
             ! grep -qE '\$\(|`' "$2"
            );then
          source $2
        else
          exmsg "$2: 敏感的文件内容"
        fi
      else
        exmsg "$2: 文件不存在"
      fi
      shift 2
      ;;
    --sysfile|-s)
      if [ -z $2 ];then
        exmsg "未指定文件"
      fi
      if [ -f $2 ];then
        sys_size=$(du -sb $2 | cut -f1) # 文件大小（字节）
        min_size=5                      # 最小 5 字节
        max_size=10240                  # 最大 10 千字节
        if ((sys_size<min_size))||((sys_size>max_size));then
          exmsg "命令提示文件内容不得 [大于 $max_size] / [小于 $min_size]"
        else
          OPENAI_SYSTEM_PROMPT="$(cat "$2")"
        fi
      fi
      shift 2
      ;;
    --help|-h)
      HELP_OPTION=(
        '--help, -h'          '显示此帮助'
        '--sysfile, -s'       '指定一个命令提示文件'
        '--config, -f'        '指定一个环境变量文件'
      )
      HELP_ENV_MUST=(
        'OPENAI_API_KEY'      '接口密钥'
        'OPENAI_API_MODEL'    '模型名称'
      )
      HELP_ENV_OPTAL=(
        'OPENAI_BASE_URL'     '起始地址（默认为 openai）'
        'OPENAI_API_TEMP'     '采样温度'
        'OPENAI_STOP_TEXTS'   '结束标语'
        'OPENAI_TEXT_NUM'     '输出文本数量'
        'OPENAI_TOP_K'        '贪婪采样阈值'
        'OPENAI_TOP_P'        '精简采样阈值'
        'OPENAI_PRESENCE'     '重复惩罚阈值'
        'OPENAI_FREQUENCY'    '频率惩罚阈值'
        'OPENAI_MAX_TOKENS'   '限制最大输出量'
      )
      echo -e "调用 OpenAI API 的终端聊天脚本

  用法: \e[32m$arg0\e[0m [ \e[33mprompt-string\e[0m | \e[33mprompt-file\e[0m | \e[33mOption\e[0m ]
  
  \e[35mOption:\e[0m
$(printf "    \e[1;34m%-27s\e[0m%-4s\n" "${HELP_OPTION[@]}")
  
  使用前需要注入环境变量
    必选值:  
$(printf "      \e[33m%-25s\e[0m%-4s\n" "${HELP_ENV_MUST[@]}")
    可选值:  
$(printf "      \e[33m%-25s\e[0m%-4s\n" "${HELP_ENV_OPTAL[@]}")

  
  未在命令参数中提问时将进入连续对话模式，可上下翻阅输入记录
  记录将保存到 \e[35m$HISTFILE\e[0m
  在连续对话模式下使用局内指令 \e[32m/help\e[0m 查看用法

  文件输入示例: \e[32mcat\e[0m \e[33mREADME.md\e[0m | \e[32m$arg0\e[0m \e[33m\"解释一下这个自述\"\e[0m\n"
      shift
      exit
      ;;
    *) # 输入提问文件或提问
      if [[ -z "$OPENAI_USER_PROMPT" ]];then
        OPENAI_USER_PROMPT="${1}"
        if [[ -f "$OPENAI_USER_PROMPT" ]];then
          OPENAI_USER_PROMPT="$(cat "$OPENAI_USER_PROMPT")"
        fi
        shift
      fi
  esac
done

########################################################################

# 设置 OpenAI API 原始地址、密钥、模型等
[[ -z ${OPENAI_API_KEY} || -z ${OPENAI_API_MODEL} ]] &&
exmsg "缺少环境变量\n需同时提供 \e[32mOPENAI_API_KEY OPENAI_API_MODEL\e[0m"
: ${OPENAI_BASE_URL:="https://api.openai.com"}
API_AUTH="Authorization: Bearer $OPENAI_API_KEY"

# 转义标准输入以确保安全
[[ -n "$STDIN_INPUT" ]] &&
STDIN_INPUT=$(printf '%s' "$STDIN_INPUT" | jq -Rsa . | sed -e 's/^"//' -e 's/"$//') &&
STDIN_INPUT="这是文件内容:\n---\n$STDIN_INPUT\n---\n"

## 筛查非法值

# 理想情况
type_of=(
  OPENAI_API_KEY,string
  OPENAI_API_MODEL,string
  OPENAI_API_TEMP,number
  OPENAI_TOP_K,number
  OPENAI_TOP_P,number
  OPENAI_PRESENCE,number
  OPENAI_FREQUENCY,number
  OPENAI_TEXT_NUM,number
  OPENAI_MAX_TOKENS,number
  OPENAI_BASE_URL,url
)

# 非法值
type_invalid=()

for i in ${type_of[@]};do
  type_value="${i%,*}"
  type_in="${i#*,}"
  type_r="$(typeval "$type_value")"
  [[ "$type_r" == "$type_in" || -z "${!type_value}" ]] ||
  type_invalid+=($type_value $type_in $type_r)
done

if ! ((${#type_invalid[@]}==0));then
  emsg "环境变量类型异常"
  printf "%-25s%-25s%-25s\n" "变量名" "类型需要" "错误指向"
  printf "%-22s%-22s%-22s\n" "${type_invalid[@]}"
  exit 1
fi

########################################################################

# 制作起始请求内容
data=$(printf '{"model": "%s", "messages": []}' "$OPENAI_API_MODEL")

# 限制 tokens 输出量
[ ! -z "$OPENAI_MAX_TOKENS" ] &&
data="$(jq --arg maxtokens "$OPENAI_MAX_TOKENS" '. += {"max_tokens": $maxtokens}' <<< "$data")"

# 指定文本数量
[ ! -z "$OPENAI_TEXT_NUM" ] &&
data="$(jq --arg num "$OPENAI_TEXT_NUM" '. += {"n": $num}' <<< "$data")"

# 温度采样
[ ! -z "$OPENAI_API_TEMP" ] &&
data="$(jq --arg temperature "$OPENAI_API_TEMP" '. += {"temperature": $temperature}' <<< "$data")"

# 贪婪采样
[ ! -z "$OPENAI_TOP_K" ] &&
data="$(jq --arg topk "$OPENAI_TOP_K" '. += {"top_k": $topk}' <<< "$data")"

# 精简采样
[ ! -z "$OPENAI_TOP_P" ] &&
data="$(jq --arg temperature "$OPENAI_TOP_P" '. += {"top_p": $topp}' <<< "$data")"

# 重复惩罚
[ ! -z "$OPENAI_PRESENCE" ] &&
data="$(jq --arg presence "$OPENAI_PRESENCE" '. += {"presence_penalty": $presence}' <<< "$data")"

# 频率惩罚
[ ! -z "$OPENAI_FREQUENCY" ] &&
data="$(jq --arg frequency "$OPENAI_FREQUENCY" '. += {"frequency_penalty": $frequency}' <<< "$data")"

# 结束标语
[ ! -z "$OPENAI_STOP_TEXTS" ] &&
array_sh2json="$(jq -n --args '$ARGS.positional' --args "${OPENAI_STOP_TEXTS[@]}")" &&
data="$(jq --argjson newarr "$array_sh2json" '. += {"stop": $newarr}' <<< "$data")"

# 提示命令
[ ! -z "$OPENAI_SYSTEM_PROMPT" ] &&
data="$(jq --arg content "$OPENAI_SYSTEM_PROMPT" '.messages += [{"role": "system", "content": $content}]' <<< "$data")"

old_data="${data}"

########################################################################

### 主要功能实现

trap 'printf "\e[?25h\r";msg "$bye\e[K\e[0m";exit' 2

# 无参数时进入连续对话模式
if [ -z "$OPENAI_USER_PROMPT" ];then
  CHAT=true
  EVAL=false
  RCSAA=false
  PRINT=true
  ERRON=false
  DEBUG=false
  SET=(
    PRINT
    ERRON
    DEBUG
    RCSAA
  )
  SET_ALIAS=(
    print-effect
    erron-effect
    debug-logs
    eval-skipask
  )
else
  CHAT=false
fi

bye="bye ~"
  
# 因 bash 脚本特性或本人能力有限，暂无法正常在输入文件时开启多轮会话
$STDIN && $CHAT && exmsg "无法处理此会话"

while true;do

  # 开启历史记录
  set -o history
    
  # 未收到/无接口返回/不在连续对话模式时不再输入聊天记录
  $CHAT && if [ ! -z "$API_RETURN" ];then
    data="$(jq --arg content "$API_RETURN" '.messages += [{"role": "assistant", "content": $content}]' <<< "$data")"
    debug "append in assistant message"
  fi

  # 用户输入
  $CHAT && read -rep "$(printf "\e[36myou:\e[0m") " OPENAI_USER_PROMPT
  ((${#OPENAI_USER_PROMPT}>0)) && echo "$OPENAI_USER_PROMPT" >> $HISTFILE

  # 空输入时直接跳过，避免浪费资源
  if [ -z "$OPENAI_USER_PROMPT" ];then
    debug "pass in enter"
    continue
  fi
  
  # 注册为指令参数
  subargv=(${OPENAI_USER_PROMPT})
  subarg0="${subargv[0]}"
  subargs=(${subargv[@]/"$OPENAI_USER_PROMPT"/})
  subargs=(${subargv[@]/'~'/$HOME})
  # 自定义局内指令
  $CHAT && case "${subarg0}" in
    /pass)      # 无操作，等价于仅回车
      debug "pass in cmd"
      continue
      ;;
    /cls)       # 清除屏幕
      clear
      debug "clear in terminal"
      continue
      ;;
    /exit|/bye) # 退出会话
      msg "$bye"
      exitcode="$[${subargs[1]}]"
      debug "exit $exitcode"
      exit $exitcode
      ;;
    /exec)      # 执行命令
      subexec="${subargs[@]:1}"
      debug "run command: $subexec"
      bash <<< "$subexec"
      CMD_RETURN=$?
      debug "status code: $CMD_RETURN"
      continue
      ;;
    /clear)     # 清空对话
      debug "clean message and file data..."
      data="${old_data}"
      unset old_filepath
      msg "对话已清空"
      continue
      ;;
    /help|/?)   # 显示帮助
      chathelp
      continue
      ;;
    /set)       # 设置
      setload ${subargs[@]:1}
      continue
      ;;
    /eval)
      OPENAI_USER_PROMPT="# 这里是bash，你只需要提供可以直接使用的命令或脚本，我将执行它，我不需要你对此做任何解释，我很清楚这样的后果。\n${subargs[@]:1}"
      EVAL=true
      ;;
    /files)      # 提交文件
      apifile ${subargs[@]:1}
      continue
      ;;
    /file)
      filepath="${subargs[1]}"
      if [ -z $filepath ];then
        erron "未输入文件"
        continue
      fi
      if [ -f $filepath ];then
        if ! filetype "$filepath";then
          erron "非法的文件类型"
          continue
        fi
        filesize=$(du -b "$filepath" | cut -f1)
        if ((filesize>102400));then
          erron "过大的文件 (${filesize}B)"
          continue
        fi
        if [ -z "$old_filepath" ];then
          STDIN_INPUT="$(cat $filepath)"
          msg "已提交文件 $filepath"
        else
          STDIN_INPUT="$(cat $filepath)"
          msg "已覆盖提交文件 $filepath => $old_filepath"
        fi
        old_filepath="$filepath"
        continue
      else
        emsg "找不到文件 $filepath"
        continue
      fi
      ;;
    /*)        # 排除未知指令
      CMD="${subarg0/'/'/}"
      if ((${#CMD}==0));then
        erron "空的指令"
      else
        erron "未知的指令 \"$CMD\""
      fi
      continue
  esac
  
  data="$(jq --arg content "$STDIN_INPUT$OPENAI_USER_PROMPT" '.messages += [{"role": "user", "content": $content}]' <<< "$data")"
  unset STDIN_INPUT

  # 构建接口请求
  CURL_GET_API="$(
    # 异步请求接口
    (
      curl -s -X POST \
        "$OPENAI_BASE_URL/chat/completions" \
        -H "$API_AUTH" \
        -H "Content-Type: application/json" \
        -d "$data" 2>/dev/null ||
      printf '{"error": {"message": "无法处理此请求"}}'
    ) &
    PROCESS=$! 
    debug "request in POST: $OPENAI_BASE_URL, model: $OPENAI_API_MODEL"
    getwait
  )"

  # 检查响应中的错误
  if ! $(jq 'has("error")' <<< "$CURL_GET_API" &>/dev/null);then
    ERROR_MSG=$(jq -r '.error.message' <<< "$CURL_GET_API" 2>/dev/null || echo "请求时发生了意料之外的错误")
    erron "$ERROR_MSG"
    unset API_RETURN
    continue
  else
    API_RETURN="$(jq --exit-status -r '.choices[0].message.content' <<< "$CURL_GET_API")"
    API_STATUS=$?
    if ((API_STATUS==1)) && [[ "$API_RETURN" == "null" ]];then
      erron "该次请求未正确被接受"
      unset API_RETURN
    else
      # 非连续对话模式下禁用逐字输出和前缀，以提高效率
      # 连续对话模式下启用逐字输出和前缀，以提升体验
      if $CHAT;then
        if $EVAL;then
          eval_code="$(
            while IFS= read -r line;do
              [[ "${line:--}" =~ ^'```' ]] || echo "$line"
            done <<<"$API_RETURN"
          )"
          if grep -qP '^\s*[\p{Han}]|^`' <<< "$eval_code";then
            cat <<< "$eval_code"
            erron "非法的命令内容"
          else
            echo -e "\e[1;34m$eval_code\e[0m"
            if ! $RCSAA;then
              msg "执行命令？[y/N] " -n
              read choose
            fi
            $RCSAA && choose=yes
            case "${choose:-no}" in
              [Yy][Ee][Ss]|[Yy])
                bash <<< "$eval_code"
                ;;
              [Nn][Oo]|[Nn])
                echo -e "\e[33m取消"
            esac
          fi
          EVAL=false
        else
          printf "\e[32mai:\e[0m "
          # 指定输出模式
          if $PRINT;then
            print <<< "$API_RETURN"
          else
            cat <<< "$API_RETURN"
          fi
        fi
      else
        cat <<< "$API_RETURN"
        break
      fi    
    fi
  fi
done